

import time, os
import torch

class DataLoader:

    def __init__(self, torchDevice:str="cuda"):
        self.torchDevice = torchDevice
        self.loadFeatureMapDict = {}


    def loadFeatureMap(self, path: str):
        # print("load: ", path)
        if path in self.loadFeatureMapDict:
            return self.loadFeatureMapDict[path]

        while not os.path.exists(path):
            time.sleep(0.5)
        fm = torch.load(path, map_location=torch.device(self.torchDevice))
        self.loadFeatureMapDict[path] = fm

        return fm.to(torch.device(self.torchDevice))

    def reload(self, l: list):
        [self.loadFeatureMap(i[1]) if i[1] in self.loadFeatureMapDict else None for i in l]


    def cleanLoadFeatureMapDict(self, l: list):
        for i in l:
            # if i[i] in self.loadFeatureMapDict:
            self.loadFeatureMap(i[1])


    def torchLoad(self, path: str):
        return torch.load(path, map_location=torch.device(self.torchDevice)) \
            .to(torch.device(self.torchDevice))


